-------------------------------------------------------------------------------
-- RTSAI - Real-Time Strategy engine for AI research
-- Copyright(C) 2006 DIMAp-UFRN, licensed under LGPL
--
-- RTSAI full capture-the-flag state machine example
-------------------------------------------------------------------------------

-------------------------------------------------------------------------------
-- Starting state
-- Behavior: if there is none or one unit seeking the flag, the unit must also
-- go to pick the flag. Otherwise, the unit must protect the flag on its lands.
-------------------------------------------------------------------------------
State {
    name = "ab_starting",  -- state after unit creation
    going_to_flag = 0,     -- num of units already going to flag

    act = function ( self )
        if self.going_to_flag < 2 then
            self.unit : broadcastMessage ( rtsai.Message.INFORM, "going to flag" )
            action.change_state.state = ab_seekingFlag : new()
        else
            action.change_state.state = ab_protectingFlag : new()
        end
        return action.change_state
    end,

    on_message = function ( self, messages )
        for i = 0, messages:size() - 1 do
            local msg = messages[i]
            if msg:get_performative() == rtsai.Message.INFORM and msg:get_content() == "going to flag" then
                self.going_to_flag = self.going_to_flag + 1
            end
        end
    end
}


-------------------------------------------------------------------------------
-- WakingUp state
-- Behavior: when a unit is woke up, it must go to pick the flag.
-------------------------------------------------------------------------------
State {
    name = "ab_wakingUp",
    act = function ( self )
        action.change_state.state = ab_seekingFlag : new()
        return action.change_state
    end
}


-------------------------------------------------------------------------------
-- SeekingFlag state
-- Behavior: Go towards the flag. However, if some other unit is frozen, it
-- comes back to help it.
-------------------------------------------------------------------------------
State {
    name = "ab_seekingFlag",  -- obrigatory field (name)

    act = function ( self )
        local player = self.unit : get_player()
        action.move.steer = self.unit : steerForSeek ( player.flag : position() )
        return action.move
    end,

    on_message = function ( self, messages )
        for i = 0, messages:size() - 1 do
            local msg = messages[i]
            if msg:get_performative() == rtsai.Message.INFORM and msg:get_content() == "going to flag" then
                self.unit : answer ( msg, rtsai.Message.INFORM, "i'm already going there" )
            end
        end
    end,

    on_unit_frozen = function ( self, event )
        if event.player == self.unit:get_player() then
            self.unit : broadcastMessage ( rtsai.Message.INFORM, "going to help unit" )
            return ab_helpingFrozenUnit : new ( { frozen = event.unit } )
        end
    end
}


-------------------------------------------------------------------------------
-- HelpingFrozenUnit state
-- Behavior: Go towards the frozen unit. When the unit is unfrozen, it goes
-- back to seek the flag.
-------------------------------------------------------------------------------
State {
    name = "ab_helpingFrozenUnit",

    act = function ( self )
        if self.frozen then
            action.move.steer = self.unit : steerForSeek ( self.frozen : position() )
            return action.move
        end
    end,

    on_unit_wakeup = function ( self, event )
        if event.unit == self.frozen then
            return ab_seekingFlag : new()
        end
    end
}


-------------------------------------------------------------------------------
-- ProtectingFlag state
-- Behavior: When some enemy unit gets on the unit lands, it attacks.
-------------------------------------------------------------------------------
State {
    name = "ab_protectingFlag",

    act = function ( self )
        local player = self.unit:get_player()
        local enemy = player.enemy
        local num_units = enemy:get_num_units()
        for i = 0, num_units - 1 do
            local e_unit = enemy:get_unit (i)
            local enemy_side = e_unit:position().z
            local enemy_land = player.flag:position().z
            if enemy_side * enemy_land < 0 then
                -- the enemy unit is not on its land, i must attack it
                action.change_state.state = ab_trackingEnemy  : new ( { enemy = e_unit } )
                return action.change_state
            end
        end
    end
}


-------------------------------------------------------------------------------
-- AttackingEnemy state
-- Behavior: Go towards an enemy in order to capture it.
-------------------------------------------------------------------------------
State {
    name = "ab_attackingEnemy",

    act = function ( self )
        if self.enemy then
            action.move.steer = self.unit : steerForSeek ( self.enemy : position() )
            return action.move
        end
    end,

    on_unit_frozen = function ( self, event )
        if event.unit == self.enemy then  -- already captured
            return ab_protectingFlag : new ()
        end
    end
}


State {
    name = "ab_trackingEnemy",
    extends = ab_attackingEnemy,

    act = function ( self )
        print ("tracking enemy")
        if self.enemy then
            action.move.steer = self.unit : steerForPursuit ( self.enemy )
            return action.move
        end
    end
}

ab_pursuitingEnemy = ab_attackingEnemy : new ( {
    act = function ( self )
        print ("pursuiting enemy")
        if self.enemy then
            action.move.steer = self.unit : steerForPursuit ( self.enemy )
            return action.move
        end
    end
    }
)
